from retrying import retry

from common.baserequest import Base_request

br = Base_request()


class EsUtils:

    @staticmethod
    def get_data(base_url: str, index: str, filter: dict, user: str, pwd: str):

        url_start = base_url + "/" + index + "/_search?scroll=1m"
        url_end = base_url + "/_search/scroll"
        res = br.post_auth(url=url_start, data=filter,
                           user=user, pwd=pwd)
        if len(res['hits']['hits']) != 0 and len(res) != 0:
            data_list = list(map(lambda x: x['_source'], res['hits']['hits']))

            yield data_list
        else:
            return
        _scroll_id = res["_scroll_id"]
        data2 = {
            "scroll": "1m",
            "scroll_id": _scroll_id
        }
        while True:
            res2 = br.post_auth(url=url_end, data=data2,
                                user=user, pwd=pwd)
            if len(res2['hits']['hits']) != 0:
                data_list = list(map(lambda x: x['_source'], res2['hits']['hits']))
                yield data_list
            else:
                break

    @staticmethod
    @retry(stop_max_attempt_number=30, wait_fixed=2000)
    def get_data_one(base_url, index, filter, user, pwd):
        url_start = base_url + "/" + index + "/_search"
        res = br.post_auth(url=url_start, data=filter,
                           user=user, pwd=pwd)
        if len(res['hits']['hits']) != 0:
            data_list = list(map(lambda x: x['_source'], res['hits']['hits']))
        else:
            data_list = []
        return data_list


if __name__ == '__main__':
    es = EsUtils()
    bass_url = "http://139.159.245.112:9208"
    index = 'allegro_listing'
    filter = {
        "size": 100,
        "query": {
            "match": {
                "OrderSourceId": 19984
            }
        }
    }
    user = "reader"
    pwd = "58Bs78E00A7D_F6"
    data_ = EsUtils.get_data(bass_url, index, filter, user, pwd)
    count = 0
    while True:
        try:
            data = next(data_)
            # print(data)
            count += len(data)
            print(count)
        except StopIteration:
            break

    print(f"共{count}")
